import numpy as np
import torch
import torch.nn as nn
from utils.graphics_utils import getProjectionMatrix
import math
import torch
from scene.gaussian_model import GaussianModel
from gaussian_renderer import render
import torchvision
from argparse import ArgumentParser
from arguments import PipelineParams


class VirtualCamera(nn.Module):
    def __init__(
        self,
        R=np.array(
            [
                [0, 0, -1],
                [0, 1, 0],
                [1, 0, 0],
            ]
        ),
        T=np.array(
            [2, 0, 0],
            dtype=np.float32,
        ),
        FoVy=math.radians(30),
        image_height=1024,
        image_width=1024,
    ):
        self.R: np.array = R
        self.T: np.array = T
        self.FoVy: float = FoVy
        self.image_height: int = image_height
        self.image_width: int = image_width
        self.update_attributes()

    def update_attributes(self):

        self.FoVx = compute_fovX(
            fovY=self.FoVy,
            height=self.image_height,
            width=self.image_width,
        )

        self.zfar = 100.0
        self.znear = 0.01

        self.world_view_transform = getViewMatrix(self.R, self.T)
        self.projection_matrix = getProjectionMatrix(
            znear=self.znear, zfar=self.zfar, fovX=self.FoVx, fovY=self.FoVy
        )
        self.full_proj_transform = self.projection_matrix @ self.world_view_transform

        self.world_view_transform = self.world_view_transform.transpose(0, 1).cuda()
        self.projection_matrix = self.projection_matrix.transpose(0, 1).cuda()
        self.full_proj_transform = self.full_proj_transform.transpose(0, 1).cuda()

        self.camera_center = torch.tensor(self.T, dtype=torch.float32).cuda()

        return

    def setLookAt(self, cam_pos, target_pos):
        cam_pos = np.array(cam_pos, dtype=np.float32)
        target_pos = np.array(target_pos, dtype=np.float32)
        y_approx = np.array([0, 1, 0], dtype=np.float32)
        self.R = LookAt(cam_pos, target_position=target_pos, y_approx=y_approx)
        self.T = cam_pos
        self.update_attributes()
        return

    def set_theta_phi_distance(self, theta, phi, distance):
        if theta > 360 or theta < 0:
            theta = theta % 360
        theta = math.radians(theta)
        phi = math.radians(phi)

        cam_pos = [
            math.cos(phi) * math.cos(theta) * distance,
            math.sin(phi) * distance,
            math.cos(phi) * math.sin(theta) * distance,
        ]

        self.setLookAt(cam_pos, target_pos=[0, 0, 0])
        return


def compute_fovX(fovY, height, width):
    w_h_ratio = float(width) / float(height)

    return math.atan(math.tan(fovY * 0.5) * w_h_ratio) * 2


def getViewMatrix(R, t):
    Rt = np.zeros((4, 4))
    Rt[:3, :3] = R.T
    Rt[:3, 3] = -R.T @ t
    Rt[3, 3] = 1.0

    return torch.tensor(Rt, dtype=torch.float32)


def LookAt(camera_position, target_position, y_approx):

    look_dir = target_position - camera_position
    z_axis = look_dir
    z_axis /= np.linalg.norm(z_axis)  # 归一化

    x_axis = np.cross(y_approx, z_axis)
    x_axis /= np.linalg.norm(x_axis)  # 归一化

    y_axis = np.cross(z_axis, x_axis)
    y_axis /= np.linalg.norm(y_axis)  # 归一化

    R = np.zeros((3, 3))
    R[:, 0] = x_axis
    R[:, 1] = y_axis
    R[:, 2] = z_axis

    return R



@torch.no_grad
def test_virtual_camera(
    fov=10,
    image_width=512,
    image_height=512,
    distance=1.1,
    theta=0,
    phi=0,
    r=0,
    g=0,
    b=0,
):
    # Set up command line argument parser
    parser = ArgumentParser(description="gaussian splatting")
    pipe = PipelineParams(parser)

    gaussians = GaussianModel(3)
    gaussians.load_ply("./models/test.ply")

    cam: VirtualCamera = VirtualCamera(
        image_width=image_width, image_height=image_height, FoVy=math.radians(fov)
    )
    # cam.setLookAt(cam_pos=[0.0, 0.3, -1.1], target_pos=[0, 0, 0])
    cam.set_theta_phi_distance(theta, phi, distance)

    bg = torch.tensor([r, g, b], dtype=torch.float32, device="cuda")
    
    with torch.no_grad():
        rendering = render(cam, gaussians, pipe, bg)["render"]
    torchvision.utils.save_image(rendering, f"./debug/img.png")

    return 

    
    


if __name__ == "__main__":

    test_virtual_camera()
